#!/usr/bin/python

import sys
import matplotlib.pyplot as plt
#import matplotlib.gridspec as gridspec
import pylab
import re
import random
from Bio import SeqIO
from Bio.SeqUtils import GC

# Usage: python /host/Users/JS/UH-work/scripts/gloeoAssemblyVerificationFigure.py GLBK_vs_phymmGloeo_pairsAndsingles.coords glbk_vs_non-gloeo-reads.coords cn.454AllContigs.coords solexa.velvet.coords atleastone.finescale.list GLBK.v1.nt
#

################################################################################
#Functions
################################################################################
def fmt(f):
    st = '{0:.4}'.format(f)
    return st

def calAvg(countlist):
    total = 0
    count = len(countlist)
    for i in countlist:
        total += i
    avg = float(total/count)
    return avg

################################################################################
#regular expressions
################################################################################
m0 = re.compile('.*0$')
m1 = re.compile('.*1$')
m2 = re.compile('.*2$')
m3 = re.compile('.*3$')
m4 = re.compile('.*4$')
m5 = re.compile('.*5$')
m6 = re.compile('.*6$')
m7 = re.compile('.*7$')
m8 = re.compile('.*8$')
m9 = re.compile('.*9$')
left = re.compile('\w+_L')
right = re.compile('\w+_R')

################################################################################
#Gloeobacter-specific reads
################################################################################
file1 = sys.argv[1]
f1 = open(file1, "rU")
fl1 = f1.readlines()

################################################################################
#Other reads
################################################################################
file2 = sys.argv[2]
f2 = open(file2, "rU")
fl2 = f2.readlines()

################################################################################
#454 contigs alignment
################################################################################
file3 = sys.argv[3]
f3 = open(file3, "rU")
fl3 = f3.readlines()

################################################################################
#Solexa contigs alignment
################################################################################
file4 = sys.argv[4]
f4 = open(file4, "rU")
fl4 = f4.readlines()

################################################################################
#phymmBL binning results
################################################################################

################################################################################
#Gloeobacter clone pairs
################################################################################
file5 = sys.argv[5]
f5 = open(file5, "rU")
fl5 = f5.readlines()

################################################################################
#Sequence file
################################################################################
seqfile = sys.argv[6]
sf = SeqIO.parse(seqfile, "fasta")
gc_values = []
gcx = []
for i, s in enumerate(sf):
    gc_values.append(GC(s.seq))
    gcx.append(i)


################################################################################
#Parsing the file1 contents
################################################################################
line = fl1[4]
l = line.split('\t')
genome_size = int(l[7])

f1_coords = {}

for i in range(1, genome_size+1):
    f1_coords[i] = 0

pairsx1 = []
pairsx2 = []
pairsy = []

singletonsx1 = []
singletonsx2 = []
singletonsy = []

for i, line in enumerate(fl1[4:]):
    l = line.split('\t')
    start = int(l[0])
    stop = int(l[1])
    readname = l[12].rstrip()
    identity = float(l[6])
    for a in range(start, stop):
        f1_coords[a] = f1_coords[a] + 1

    if left.match(readname) or right.match(readname):
        pairsx1.append(start)
        pairsx2.append(stop)
        if m0.match(str(start)):
            pairsy.append(11.2)
        elif m1.match(str(start)):
            pairsy.append(11.4)
        elif m2.match(str(start)):
            pairsy.append(11.6)
        elif m3.match(str(start)):
            pairsy.append(11.8)
        elif m4.match(str(start)):
            pairsy.append(12)
        elif m5.match(str(start)):
            pairsy.append(12.2)
        elif m6.match(str(start)):
            pairsy.append(12.4)
        elif m7.match(str(start)):
            pairsy.append(12.6)
        elif m8.match(str(start)):
            pairsy.append(12.8)
        elif m9.match(str(start)):
            pairsy.append(13)
    else:
        singletonsx1.append(start)
        singletonsx2.append(stop)
        if m0.match(str(start)):
            singletonsy.append(11.2)
        elif m1.match(str(start)):
            singletonsy.append(11.4)
        elif m2.match(str(start)):
            singletonsy.append(11.6)
        elif m3.match(str(start)):
            singletonsy.append(11.8)
        elif m4.match(str(start)):
            singletonsy.append(12)
        elif m5.match(str(start)):
            singletonsy.append(12.2)
        elif m6.match(str(start)):
            singletonsy.append(12.4)
        elif m7.match(str(start)):
            singletonsy.append(12.6)
        elif m8.match(str(start)):
            singletonsy.append(12.8)
        elif m9.match(str(start)):
            singletonsy.append(13)

sxa = [singletonsx1, singletonsx2]
sxb = [singletonsy, singletonsy]

#sxa = [pairsx1, pairsx2]
#sxb = [pairsy, pairsy]

cov1 = []

for k, v in f1_coords.iteritems():
    cov1.append(v)

w = 1000
depthx1 = []
depthy1 = []
while w < len(cov1):
    average = calAvg(cov1[w-1000:w])
    depthx1.append(w-1000+500)
    depthy1.append(average)
    w += 1000

################################################################################
#Parsing the file2 contents
################################################################################

f2_coords = {}
for i in range(1, genome_size+1):
    f2_coords[i] = 0

for j, line in enumerate(fl2[4:]):
    l = line.split('\t')
    start = int(l[0])
    stop = int(l[1])
    identity = float(l[6])
    for a in range(start, stop):
        f2_coords[a] = f2_coords[a] + 1

cov2 = []

for k, v in f2_coords.iteritems():
    cov2.append(v)

x = 1000
depthx2 = []
depthy2 = []
while x < len(cov2):
    average = calAvg(cov2[x-1000:x])
    depthx2.append(x+500)
    depthy2.append(average)
    x += 1000

#For line representing the assembled genome
gx = [1,genome_size]
gy = [1,1]

xcoords = range(1, genome_size+1)
#ycoords = cov

################################################################################
#Parsing the file3 contents
#For 454 contigs assembled with Newbler
################################################################################

x1 = []
x2 = []
y454 = []

for i, line in enumerate(fl3[4:]):
    l = line.split('\t')
    start = int(l[0])
    stop = int(l[1])
    x1.append(start)
    x2.append(stop)
    if (i+1-1)%2 == 0:
        y454.append(6)
    else:
        y454.append(5.5)
a = [x1, x2]
b = [y454, y454]

################################################################################
#Parsing the file4 contents
#For Solex contigs assembled with Velvet
################################################################################
x3 = []
x4 = []
ys = []

for i, line in enumerate(fl4[4:]):
    l = line.split('\t')
    start = int(l[0])
    stop = int(l[1])
    x3.append(start)
    x4.append(stop)
    if (i+1-1)%2 == 0:
        ys.append(4)
    else:
        ys.append(3.5)
c = [x3, x4]
d = [ys, ys]

################################################################################
#Parsing the file5 contents
#Gloeobacter clone pairs
################################################################################
gpx1 = []
gpx2 = []
gpy = []

for i, line in enumerate(fl5):
    l = line.split('\t')
    start = int(l[2].split(' ')[0])
    stop = int(l[2].split(' ')[2])
    gpx1.append(start)
    gpx2.append(stop)
    if m0.match(str(i)):
        gpy.append(7.2)
    elif m1.match(str(i)):
        gpy.append(7.4)
    elif m2.match(str(i)):
        gpy.append(7.6)
    elif m3.match(str(i)):
        gpy.append(7.8)
    elif m4.match(str(i)):
        gpy.append(8)
    elif m5.match(str(i)):
        gpy.append(8.2)
    elif m6.match(str(i)):
        gpy.append(8.4)
    elif m7.match(str(i)):
        gpy.append(8.6)
    elif m8.match(str(i)):
        gpy.append(8.8)
    elif m9.match(str(i)):
        gpy.append(9)

gpa = [gpx1, gpx2]
gpb = [gpy, gpy]

################################################################################
#Close the file handlers
################################################################################
f1.close()
f2.close()
f3.close()
f4.close()

################################################################################
#Print notice..
################################################################################
print "# of records in xcoords: ", len(xcoords)
print "# of records in cov1: ", len(cov1)


################################################################################
# Start plotting
################################################################################
fig = plt.figure(1, figsize=(15,8))
#plt.subplots_adjust(wspace=4.0)
################################################################################
# First subplot
################################################################################
#plt.subplot(311)
#plt.axis([0, genome_size, 0, 20])
#plt.plot(gpa, gpb, color='purple', linestyle='-')
#plt.title('Gloeobacter assembly verification')
#plt.grid(True)

################################################################################
# Second subplot
################################################################################
plt.subplot(311)
#gs = gridspec.GridSpec(3, 3)
#plt.subplot(gs[0:1, 0:2])
plt.title('Assembly verification')
plt.plot(gx, gy, color='#AF7817', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)
plt.plot(sxa, sxb, color='#333366', linestyle='-')
plt.plot(a, b, color='purple', linestyle='-', lw=2.0)
plt.plot(c, d, color='red', linestyle='-', lw=2.0)
plt.plot(gpa, gpb, color='#ADA96E', linestyle='-')
plt.annotate('Singleton reads', xy=(0.8, 0.95), xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=9)
plt.annotate('Mate pairs (only showing one pair for a given segment)', xy=(0.8, 0.68), xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=9)
plt.annotate('454 contigs (Newbler)', xy=(0.8, 0.47), xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=9)
plt.annotate('Solexa contigs (Velvet)', xy=(0.8, 0.20), xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=9)

plt.axis([0, genome_size, 0, 14])
frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)
plt.grid(False)

################################################################################
# Third subplot
################################################################################

plt.subplot(312)
#plt.subplot(gs[2, :])
#plt.fill_between(xcoords, cov1, facecolor='#660066', alpha=0.4)
#plt.fill_between(xcoords, cov2, facecolor='#33FF33', alpha=0.4)
plt.fill_between(depthx2, depthy2, facecolor='#33FF33', alpha=0.4)
plt.fill_between(depthx1, depthy1, facecolor='#660066', alpha=0.4)
#plt.annotate('test', xy=(40000, 40), xycoords='data', xytext=(30, 30), textcoords='offset points', arrowprops=dict(arrowstyle="->"))
plt.annotate('Green = other organisms', xy=(0.2, 0.80), xycoords='axes fraction', horizontalalignment='center', verticalalignment='center')
plt.annotate('Purple = Gloeobacter', xy=(0.2, 0.60), xycoords='axes fraction', horizontalalignment='center', verticalalignment='center')
plt.axvspan(2726297, 2731211, facecolor='blue', alpha=0.4)
plt.annotate('rRNA operon', xy=(2731211, 40), xycoords='data', xytext=(30, 0), 
    textcoords='offset points', arrowprops=dict(arrowstyle="->"))
#plt.axis([0, genome_size, 0, max(cov1)+2])
plt.axis([0, genome_size, 0, max(depthy1)+2])
plt.xlabel('Genome position (bp)')
plt.ylabel('Counts of binned reads (per 1000bp)')
plt.grid(True)

plt.subplot(313)
plt.plot(gcx, gc_values, color='blue', linestyle='-', alpha=0.4)
plt.xlabel("Genes")
plt.ylabel("G+C%")

plt.show()

#plt.savefig(outfile, format='pdf')
